from vision_utils import *
from traceback import print_exc
import serial
import time
import json

# from cali_util import calculate_target_position

cap = cv2.VideoCapture("/dev/video20")
driver = serial.Serial("/dev/ttyS8", 115200, timeout=1)


def reset(p=False):
    if not p:
        driver.write("R".encode("ascii"))
        logger.info("Reset")
    else:
        driver.write("X120.0,Y100.0E".encode("ascii"))
        logger.info("Reset to 120,100")


def move_to(posi):
    message = f"X{posi[0]},Y{posi[1]}E"
    logger.info(f"Sending {message}")
    driver.write(message.encode("ascii"))

    for i in range(5):
        message = driver.read()
        if message == b"Z":
            logger.info(f"Received {message}, Can not reach area")
            return False
        time.sleep(0.1)

    return True


def get_mark_for_cali(frame):
    # 检测ArUco标记
    markers, ids, _ = aruco.detectMarkers(frame, aruco_dict, parameters=parameters)

    # return corners, ids

    ori_posi = None
    ori_posi_detailed = None
    board_mark_up = None
    board_mark_down = None

    if markers is not None:
        if markers[0] is not None:
            for mark_index in range(len(markers)):
                center_posi = np.mean(markers[mark_index][0], axis=0)
                mark_id = ids[mark_index][0]
                if mark_id == 4:
                    ori_posi = center_posi
                    ori_posi_detailed = markers[mark_index]
                elif mark_id == 0:
                    board_mark_up = center_posi
                elif mark_id == 9:
                    board_mark_down = center_posi

                # cv2.circle(
                #     frame,
                #     (int(center_posi[0]), int(center_posi[1])),
                #     5,
                #     (0, 0, 255),
                #     10,
                # )

        else:
            logger.warning("No marks, wait for next frame")
            raise ValueError("No marks, wait for next frame")
    else:
        raise ValueError("No marks, wait for next frame")

    if ori_posi is None:
        logger.warning("No marks, wait for next frame")
        raise ValueError("No marks, wait for next frame")

    return (
        {"center": ori_posi, "detail": ori_posi_detailed},
        board_mark_up,
        board_mark_down,
    )


result = []
reset()
input("Press enter when ready >>>")


ret, frame = cap.read()
frame = cv2.flip(frame, -1)
dbg_img = frame.copy()
get_main_status(cap)

print(ori_scale)

input("Press enter when ready >>>")

# {'center', 'valid', 'offset'}
dbg_cmd = ""

for x in range(8):
    for y in range(6):
        target = (x * 80 + 40, y * 80 + 40)
        if move_to(get_real_posi(target, use_cali=False)):
            dbg_cmd += input("Please place aruco tag >>>")
            reset()
            dbg_cmd += input("Press enter when reset done >>>")

            if "n" in dbg_cmd:
                reset()
                input("Press enter when reset done >>>")
                result.append(
                    {
                        "center": target,
                        "valid": False,
                    }
                )
                cv2.circle(frame, target, 5, (255, 0, 0), 10)

                cv2.imshow("frame", frame)
                cv2.waitKey(1)

                dbg_cmd = ""

                input("Next sample >>>")
                continue
            if dbg_cmd != "":
                reset(True)

            while True:
                for _ in range(10):
                    ret, dbg_img = cap.read()
                    dbg_img = cv2.flip(dbg_img, -1)

                if not ret:
                    logger.error("No frame in camera")
                    cap.release()
                    cap = cv2.VideoCapture("/dev/video1")
                    continue

                try:
                    ori_posi, board_mark_up, board_mark_down = get_mark_for_cali(
                        dbg_img
                    )
                    # draw markers
                except:
                    if input("give up sample ? >>>") == "y":
                        result.append(
                            {
                                "center": target,
                                "valid": False,
                            }
                        )
                        cv2.circle(frame, target, 5, (255, 0, 0), 10)
                        break
                    print_exc()
                    continue

                result.append(
                    {
                        "center": target,
                        "valid": True,
                        "offset": (
                            -(target[0] - ori_posi["center"][0]),
                            -(target[1] - ori_posi["center"][1]),
                        ),
                    }
                )

                print(result[-1])

                cv2.circle(frame, target, 5, (0, 0, 255), 10)
                cv2.circle(
                    frame,
                    tuple(map(lambda x: int(x), tuple(ori_posi["center"]))),
                    5,
                    (0, 255, 255),
                    10,
                )

                break

            if dbg_cmd != "":
                reset()

            dbg_cmd = ""

        else:
            reset()
            input("Press enter when reset done >>>")
            result.append(
                {
                    "center": target,
                    "valid": False,
                }
            )
            cv2.circle(frame, target, 5, (255, 0, 0), 10)

        cv2.imshow("frame", frame)
        cv2.waitKey(1)

        input("Next sample >>>")
        # cv2.circle(frame, target, 5, (255, 0, 255), 10)
# exit()
print(result)
with open("cali.json", "w") as f:
    json.dump(result, f)

exit()


while True:
    ret, frame = cap.read()
    dbg_img = frame.copy()

    if not ret:
        logger.error("No frame in camera")
        cap.release()
        cap = cv2.VideoCapture("/dev/video1")
        continue

    img_size = frame.shape
    # logger.debug(f"Frame size: {img_size}")
    # frame = undistort_frame(frame)
    # frame = cv2.flip(frame, 1)
    try:
        ori_posi, board_mark_up, board_mark_down = get_mark_for_cali(frame)
        # draw markers
        cv2.circle(
            frame,
            tuple(map(lambda x: int(x), tuple(board_mark_up))),
            5,
            (0, 0, 255),
            10,
        )
        cv2.circle(
            frame,
            tuple(map(lambda x: int(x), tuple(ori_posi["center"]))),
            5,
            (0, 0, 255),
            10,
        )
        cv2.circle(
            frame,
            tuple(map(lambda x: int(x), tuple(board_mark_down))),
            5,
            (0, 0, 255),
            10,
        )
        # get center of up and down
        center0 = tuple(
            map(
                lambda x: int(x),
                tuple(
                    (
                        (board_mark_up[0] + board_mark_down[0]) / 2,
                        (board_mark_up[1] + board_mark_down[1]) / 2,
                    )
                ),
            )
        )
        cv2.circle(frame, center0, 5, (0, 255, 255), 10)
    except:
        print_exc()
        continue

    for x in range(8):
        for y in range(6):
            target = (x * 80 + 40, y * 80 + 40)
            cv2.circle(frame, target, 5, (255, 0, 255), 10)

    cv2.imshow("frame", frame)

    cv2.waitKey(1)
